args=commandArgs(T)

TCRDirListi=args[1] 
outTCRFile=args[2] 
flag=args[3]  #1,all; 2:1A1B,1A2B,2A1B 3:only 1A0B or 0A1B

#output two files:
#TCRTable and TCRTable.StandTCR.Bar

#TCRDirListi:
#P4_Ptumor_	/data2/users/dongxingjun/Human_SC_AnaData/P4/Ptumor/TCR/outs

#finish in 2020.2.8
#from cellrange's TCR analysis data get TCR Table

readList <- function(LTable){
  a <- read.table(LTable,sep = "\t")
  sampleList  <- as.vector(a$V1)
  b <- as.vector(a$V2)
  names(b) <- as.vector(a$V1)
  return(b)
}


readTCR <- function(data.dir=NULL,CellType=NULL)
{
  #get all TCR info
  
  contig.loc <- paste0(data.dir,"/filtered_contig_annotations.csv")
  contig <- read.csv(contig.loc)
  all.cell <- as.vector(contig[,1])
  
  contig[,1]=all.cell
  estimated.cell <- unique(all.cell)
  
  productiveA.cell <- unique(all.cell[which(contig$productive=="True" & contig$chain=="TRA")])
  productiveB.cell <- unique(all.cell[which(contig$productive=="True" & contig$chain=="TRB")])
  productiveAB.cell <- union(productiveA.cell,productiveB.cell)

  id <- apply(contig[,c("v_gene","d_gene","j_gene","cdr3")],1, function(x)paste(x,collapse = ":")
  )
  clonotype2cell <- cbind(id,contig[,c("barcode","chain","raw_clonotype_id")])[which(contig$productive=="True"),]
  clonotype2cell <- clonotype2cell[which(clonotype2cell$barcode %in% productiveAB.cell),]

  #get final table
  all.TCR.srt=lapply(unique(clonotype2cell$barcode),function(x)
  {
    idx=which(clonotype2cell$barcode==x)
    TCR=clonotype2cell[idx,]
    
    TRA=TCR[which(TCR$chain=="TRA"),"id"]
    TRA.count=length(TRA)
    TRA=paste(TRA,collapse = "|")
    
    TRB=TCR[which(TCR$chain=="TRB"),"id"]
    TRB.count=length(TRB)
    TRB=paste(TRB,collapse = "|")
    
    TCR.id=paste(TRA,TRB,sep=";")
    Tcount=TRA.count+TRB.count
    
    TCR.stat = paste(TRA.count,"A",TRB.count,"B",sep="")
    
    c(TCR.id,TRA.count,TRB.count,Tcount,TCR.stat)
    
  })

  all.TCR.srt=do.call(rbind,all.TCR.srt)
  barcode <- unique(clonotype2cell$barcode)
  barcode <- paste(CellType,barcode,sep="")
  rownames(all.TCR.srt) = barcode
  
  finalTable = as.data.frame(all.TCR.srt)
  finalTable = cbind(barcode,finalTable)
  TableColName <- c("barcode","TCRid","ACount","BCount","TotCount","TCRstat")
  colnames(finalTable) <- TableColName

  return(finalTable)
}

#main
a <- readList(TCRDirListi)
cellTypeList <- names(a)

#for every sample
finalTable <- NULL
for(i in seq(1,length(cellTypeList),1))
{
  celltype <- cellTypeList[i]
  dataDir <- as.character(a[i])

  print(celltype)
  print(dataDir)
  tempTable <- readTCR(data.dir=dataDir,CellType=celltype)      
  finalTable <- rbind(finalTable,tempTable)
}

#filter
if(flag == 1)
{
  #all kinds of TCR (XAXB)
  print("all kinds of TCR barcode:")
  finalTable <- finalTable
}else if(flag == 2)
{
  print("TCR:1A1B,2A1B,1A2B,2A2B:")
  #1A1B,2A1B,1A2B
  finalTable <- finalTable[which(finalTable$ACount %in% c(1,2) & finalTable$BCount %in% c(1,2) & finalTable$TotCount %in% c(2,3,4)),]
}else if(flag == 3)
{
  print("TCR:1A0B,0A1B:")
  #only 1A0B or 0A1B
  finalTable <- finalTable[which(finalTable$ACount %in% c(0,1) & finalTable$BCount %in% c(0,1) & finalTable$TotCount %in% c(1)),]
}


#output final table
write.table(finalTable,file = outTCRFile,row.names = F,col.names = T,sep="\t",quote=F)
outBar <- paste(outTCRFile,".StandTCR.Bar",sep="")
print(outBar)
outBarTCRStat <- c("1A1B","1A2B","2A1B","2A2B")
writeLines(as.vector(finalTable[which(finalTable$TCRstat %in% outBarTCRStat),]$barcode),con=outBar,sep="\n")

